import os
from annoy import AnnoyIndex
from app import conf

DIM = 2
t = AnnoyIndex(DIM, metric="angular")
filename = os.path.join(conf.DATA_ROOT, 'test.ann')


def save():
	vecs = [[1, 2], [3, 4], [1, 3], [1, 4], [1, 5], [2, 3],[2, 4]]
	
	for i, v in enumerate(vecs):
		t.add_item(i, v)
	t.build(3)
	t.save(filename)


def query(key: int, topK=3):
	t.load(filename)
	ids, distances = t.get_nns_by_item(key, topK+1, search_k=-1, include_distances=True)
	for i in range(len(ids)):
		if ids[i] == key:
			del ids[i]
			del distances[i]
			break
	return ids, distances


if __name__ == '__main__':
	# save()
	print(query(0))
